package 力扣._00周赛._426_周赛.问题3;

import java.util.*;

class Solution2 {
    public int[] maxTargetNodes(int[][] edges1, int[][] edges2, int k) {
        int m = edges1.length + 1, n = edges2.length + 1;

        // 构建两棵树的邻接表
        List<List<Integer>> adj1 = buildAdjList(edges1, m);
        List<List<Integer>> adj2 = buildAdjList(edges2, n);

        // 预处理两棵树中每个节点在深度 k 内的可达节点数
        int[][] reachCount1 = preprocessReachCount(adj1, m, k);
        // 图2只需要计算到 k-1
        int[][] reachCount2 = preprocessReachCount(adj2, n, k-1);

        // 预先计算 reachCount1 和 reachCount2 的累积和
        int[][] cumulativeSum1 = new int[m][k + 1];
        int[][] cumulativeSum2 = new int[n][k];
        for (int i = 0; i < m; i++) {
            cumulativeSum1[i][0] = reachCount1[i][0];
            for (int d = 1; d <= k; d++) {
                cumulativeSum1[i][d] = cumulativeSum1[i][d - 1] + reachCount1[i][d];
            }
        }
        for (int j = 0; j < n; j++) {
            cumulativeSum2[j][0] = reachCount2[j][0];
            for (int d = 1; d < k; d++) {
                cumulativeSum2[j][d] = cumulativeSum2[j][d - 1] + reachCount2[j][d];
            }
        }

        // 遍历第一棵树的每个节点 i，模拟与第二棵树中每个节点 j 连接
        int[] res = new int[m];
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                // 合并两棵树的 reachCount
                int count = cumulativeSum1[i][k];
                if (k > 0) {
                    count += cumulativeSum2[j][k - 1];
                }
                res[i] = Math.max(res[i], count);
            }
        }
        return res;
    }

    // 构建邻接表
    private List<List<Integer>> buildAdjList(int[][] edges, int size) {
        List<List<Integer>> adj = new ArrayList<>(size);
        for (int i = 0; i < size; i++) {
            adj.add(new ArrayList<>());
        }
        for (int[] edge : edges) {
            adj.get(edge[0]).add(edge[1]);
            adj.get(edge[1]).add(edge[0]);
        }
        return adj;
    }

    // 预处理可达节点数
    private int[][] preprocessReachCount(List<List<Integer>> adj, int size, int k) {
        if (k == -1) {
            int[][] ints = new int[size][1];
            for (int i = 0; i < size; i++) {
                ints[i][0] = 1;
            }
            return ints;
        }
        int[][] reachCount = new int[size][k + 1];
        for (int i = 0; i < size; i++) {
            bfs(adj, i, size, k, reachCount);
        }
        return reachCount;
    }

    // 使用 BFS 计算每个节点在深度 k 内的可达节点数
    private void bfs(List<List<Integer>> adj, int start, int size, int k, int[][] reachCount) {
        Queue<Integer> queue = new LinkedList<>();
        boolean[] visited = new boolean[size];
        queue.add(start);
        visited[start] = true;
        reachCount[start][0] = 1; // 自己

        int level = 0;
        while (!queue.isEmpty() && level < k) {
            int levelSize = queue.size();
            for (int i = 0; i < levelSize; i++) {
                int node = queue.poll();
                for (int neighbor : adj.get(node)) {
                    if (!visited[neighbor]) {
                        visited[neighbor] = true;
                        queue.add(neighbor);
                        reachCount[start][level + 1]++;
                    }
                }
            }
            level++;
        }
    }

    public static void main(String[] args) {
        int[][] edges1 = new int[][]{{0, 1}, {0, 2}, {2, 3}, {2, 4}};
        int[][] edges2 = new int[][]{{0, 1}, {0, 2}, {0, 3}, {2, 7}, {1, 4}, {4, 5}, {4, 6}};
        int k = 2;
        Solution2 solution = new Solution2();
        int[] res = solution.maxTargetNodes(edges1, edges2, k);
        for (int i : res) {
            System.out.println(i);
        }
    }
}